import yaml
from pathlib import Path
from ultralytics import YOLO


# 推理模型
def inference():
    with open("config.yaml", "r", encoding="utf-8") as file:
        data = yaml.safe_load(file)
    print(data)
    input_path = Path(data["inference"]["input_filepath"])
    output_path = Path(data["inference"]["output_filepath"])
    model = YOLO(data["inference"]["checkpoint"])
    for img_path in input_path.iterdir():
        results = model(img_path, imgsz=1024, device="cpu", conf=0.25)
        for result in results:
            result.save(filename=output_path / f"{img_path.stem}.jpg", conf=False)


if __name__ == "__main__":
    inference()
